#!/usr/bin/env python
# coding=utf-8
# @Time    : 2024/1/31 19:13
# @Software: PyCharm
from fastapi import FastAPI
from fastapi.responses import StreamingResponse
from fastapi.middleware.cors import CORSMiddleware
import time
import uvicorn
from client import run_agent
import asyncio

app = FastAPI()

# 添加CORS中间件配置
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],  # 允许所有来源
    allow_credentials=True,
    allow_methods=["*"],  # 允许所有方法
    allow_headers=["*"],  # 允许所有头
)

# async def generate_data():
#     for i in range(1, 11):
#         time.sleep(1)  # 模拟每秒生成一个块的耗时操作
#         yield f"FASTAPI Chunk {i}\n"
# agent = asyncio.run(run_agent())


async def generate_data():
    for i in range(1, 11):
        time.sleep(1)  # 模拟每秒生成一个块的耗时操作
        yield f"FASTAPI Chunk {i}\n"


@app.get("/stream")
async def stream_data():

    # 添加CORS相关响应头
    return StreamingResponse(
        generate_data(),
        media_type="text/plain",
        headers={"Access-Control-Allow-Origin": "*", "Cache-Control": "no-cache"},
    )


# 添加启动代码（可选）
if __name__ == "__main__":
    uvicorn.run(app, host="0.0.0.0", port=8000)
    # generate_data()
